#include <iostream>
#include <eigen3/Eigen/Sparse>
#include <vector>
#include <functional>
#include <ctime>
#include <eigen3/Eigen/SparseLU>
#include <cmath>
typedef unsigned int Dimension;
typedef double Real;
typedef std::vector<Real> Vector;
typedef Eigen::VectorXd VectorXd;
typedef Eigen::MatrixXd MatrixXd;
typedef Eigen::Triplet<Real> Tri;
typedef Eigen::SparseMatrix<Real> SpM;
typedef std::vector<SpM> SpMbox;
typedef std::function<double(double *)> Function;

enum Restriction
{
    full_weighting = 1,
    injection = 2
};

enum Interpolation
{
    linear = 1,
    quadratic = 2
};

/**
 * @brief 多重网格求解器。是一个蕴含了TG,VC与FMG的共性的纯虚基类。
 */
class MultigridSolver
{
protected:
    /**
    * @brief 待求解方程组系数矩阵;
    */
    MatrixXd _A_origin;
    /**
     * @brief 用于存储基本系数矩阵集合；
    */
    SpMbox _A_basic;
    /**
     * @brief 待求解线性方程组右端项向量；
     */
    VectorXd _Rhs_origin;
    /**
     * @brief 计算过程中用于赋值和保存的临时右端向量；
     */
    VectorXd _Rhs_temp;

    /**
     * @brief 
     */
    VectorXd _u_temp;
    /**
     * @brief 用于存储多重网格计算过程中，所产生的各线性方程组右端项向量；
     */
    std::vector<VectorXd> _Rhs;
    /**
     * @brief 用于存储多重网格计算过程中，所产生的线性方程组的解或者参量方程的解；
     */
    std::vector<VectorXd> _u;
    /**
     * @brief 用于记录边界条件函数；
     */
    Function _BCF;
    /**
     * @brief 用于记录当前多重网格求解器所在层数；
     */
    int _nowlevel = 0;
    /**
     * @brief 设置迭代终止条件之一:||Au - f||L2 <该值时计算终止；
     */
    Real _Tolerance = 1e-6;
    /**
     * @brief 设置迭代终止条件之一:迭代最大步数；
     */
    long _MaxIternum = 20;
    /**
     * @brief 用于存放所有插值算子，在initialize函数中进行初始化;
     */
    SpMbox _Restric;
    /**
     * @brief 用于存放所有限制算子，在initialize函数中进行初始化;
     */
    SpMbox _Intp;
    /**
     * @brief 用于记录最大下降层数；
     */
    int _maxlevel = 5;

    int _numberofgrid = 128;

    enum Restriction _Restric_method = full_weighting;

    enum Interpolation _Intp_method = linear;

public:
    /**
     * @brief Set the Tolerance 
     * 
     * @param tolerance 最大容许误差;
     */
    void setTolerance(Real tolerance) { _Tolerance = tolerance; };
    /**
     * @brief Set the Maxlevel object
     * @param  maxlevel        最大下降层数;
     */
    void setMaxlevel(int maxlevel) { _maxlevel = maxlevel; };
    /**
     * @brief 上升一层；
     */
    void UpLevel() { _nowlevel++; };
    /**
     * @brief 下降一层；
     */
    void DownLevel() { _nowlevel--; };

    /**
     * @brief 传统迭代求解器。这里选择了w=2/3的加权Jacobi.
     * @param  _A               系数矩阵
     * @param  _f               右端项
     * @param  _v               迭代初值
     * @param  _k_times          迭代次数
     * @return VectorXd 
     */
    VectorXd Relax_solver(SpM _A, VectorXd _f, VectorXd _v, int _k_times);
    /**
     * @brief 生成原始右端项。
     * @param  n                1/n是最细的网格宽度
     * @param  _bcf             边界条件函数
     */

    /**
     * @brief 一维特供
     * @param  n                My Param doc
     * @param  _F               My Param doc
     * @param  _start           My Param doc
     * @param  _end             My Param doc
     */
    void GenerateRhs_origin(int n, Function _F, double _start, double _end);
    /**
     * @brief 生成插值和限制算子；
     */
    void GenerateP(int n);
    /**
     * @brief 生成基本系数矩阵A
     * @param  n                1/n是最细的网格宽度
     */
    void GenerateA_basic(int n);

    void Initialize(int n, Function _F, double _start, double _end);
    VectorXd Vcycle();
    VectorXd Jacobi();
};



void MultigridSolver::Initialize(int n, Function _F, double _start, double _end)
{
    std::cout << "Initializing..." << std::endl;
    _numberofgrid = n;
    _nowlevel = 0;
    GenerateRhs_origin(n, _F, _start, _end);

    GenerateP(n);

    GenerateA_basic(n);

    _Rhs.reserve(_maxlevel + 1);

    _u.reserve(_maxlevel + 1);

    int _power_record = 1;
    for (int i = 0; i < _maxlevel + 1; i++)
    {
        _Rhs.push_back(VectorXd::Zero(n / _power_record - 1));
        _u.push_back(VectorXd::Zero(n / _power_record - 1));
        _power_record *= 2;
    };
    _Rhs[0] = _Rhs_origin;
    _u[0] = VectorXd::Zero(_Rhs[0].rows());

    std::cout << "Initialized!" << std::endl;
};

VectorXd MultigridSolver::Jacobi()
{
    return Relax_solver(_A_basic[_nowlevel]*_numberofgrid*_numberofgrid,_Rhs[_nowlevel],_u[_nowlevel],5000);
};

VectorXd MultigridSolver::Vcycle()
{   
    double _h_temp=1.0/_numberofgrid;
    // std::cout << "begin with" << std::endl;
    // std::cout << "_A:" << _A_basic[_nowlevel].rows() << std::endl;
    // std::cout << "_u:" << _u[_nowlevel].rows() << std::endl;
    // std::cout << "_Rhs" << _Rhs[_nowlevel].rows() << std::endl;
    _u[_nowlevel] = Relax_solver(_A_basic[_nowlevel]/(_h_temp*_h_temp), _Rhs[_nowlevel], _u[_nowlevel], 1);

    if (_nowlevel < _maxlevel)
    {

        UpLevel();
        _h_temp*=2;
        // std::cout << "_nowlevel up up,is " << _nowlevel << std::endl;
        // std::cout << "_RHs:" << _Rhs[_nowlevel].rows() << std::endl;

        // std::cout << "_Restric:" << _Restric[_nowlevel].cols() << std::endl;

        _Rhs[_nowlevel] = _Restric[_nowlevel - 1] * (_Rhs[_nowlevel - 1] - _A_basic[_nowlevel - 1] * _u[_nowlevel - 1]);
        // std::cout << "_Restric:" << _Restric[_nowlevel - 1].cols() << std::endl;
        // std::cout << ":(" << std::endl;
        
        _u[_nowlevel] = Vcycle();
        DownLevel();
        _h_temp/=2;
        // std::cout << "_nowlevel down down,is " << _nowlevel << std::endl;
        // std::cout << "_u:" << _u[_nowlevel].rows() << std::endl;
        // std::cout << "_Intp:" << _Intp[_nowlevel ].cols() << std::endl;
        // std::cout << "_u(past):" << _u[_nowlevel + 1].rows() << std::endl;
        _u_temp=_u[_nowlevel];
        _u[_nowlevel] = _u[_nowlevel] + _Intp[_nowlevel] * _u[_nowlevel + 1];
        // std::cout<<"Show me difference"<<std::endl;
        // std::cout<< _u_temp.norm()-_u[_nowlevel].norm() << std::endl;

        // std::cout << "end" << std::endl;
    }
    // else if(_nowlevel==_maxlevel)
    // {   

    //     return _u[_nowlevel].SparseLu().solve(_Rhs[_nowlevel]);
    // };
    // else
    // {

	// Eigen::LeastSquaresConjugateGradient<Eigen::SparseMatrix<double> > Solver_sparse;
 
    // // 设置迭代精度
	// Solver_sparse.setTolerance(0.0001);
	// Solver_sparse.compute(_A_basic[_nowlevel]);
 
    // //x1_sparse 即为解
	// _u[_nowlevel] = Solver_sparse.solve(_Rhs[_nowlevel]);
    // };

    // if(_nowlevel<_maxlevel)
    _u[_nowlevel] = Relax_solver(_A_basic[_nowlevel]/(_h_temp*_h_temp), _Rhs[_nowlevel], _u[_nowlevel], 1);

    return _u[_nowlevel];
};

void MultigridSolver::GenerateRhs_origin(int n, Function _F, double _start, double _end)
{
    double a[1] = {0};

    _Rhs_origin.resize(n - 1);

    double h=1.0/n;
    for (int i = 0; i < n - 2; i++)
    {
        a[0] = 1.0 * (i + 1) / n;
        _Rhs_origin(i) = _F(a);
    };
    _Rhs_origin(0) += n*n*_start;
    _Rhs_origin(n - 2) += n*n*_end;
    // std::cout<<"_RHS_ORIGIN:"<<std::endl;
    // std::cout<<_Rhs_origin<<std::endl;
};

VectorXd MultigridSolver::Relax_solver(SpM _A, VectorXd _f, VectorXd _v, int _k_times)
{
    double w = 2.0/3;

    int _width = _A.rows();
    SpM D_inv, T;
    VectorXd _v_temp = _v;
    D_inv = _A.diagonal().asDiagonal().inverse();
    T.resize(_width, _width);
    T.setIdentity();
    T = T - w * (D_inv * _A);

    /**
     * @brief test
     */
    // std::cout<<"_v: "<< _v.rows()<<std::endl;
    // std::cout<<"iteration: "<< T.rows() <<std::endl;
// std::cout<<"Show you"<<std::endl;
// std::cout<<T<<std::endl;
    for (int i = 0; i < _k_times; i++)
        _v_temp = T * _v_temp + w * (D_inv * _f);
    return _v_temp;
};

void MultigridSolver::GenerateA_basic(int n)
{
    int _maxdepth = _maxlevel + 1;
    int _width = n;
    std::vector<Tri> tripletList;
    tripletList.reserve(3 * n - 5);
    tripletList.push_back(Tri(0, 0, 2));
    tripletList.push_back(Tri(0, 1, -1));
    for (int i = 1; i < n - 2; i++)
    {
        tripletList.push_back(Tri(i, i - 1, -1));
        tripletList.push_back(Tri(i, i, 2));
        tripletList.push_back(Tri(i, i + 1, -1));
    };
    tripletList.push_back(Tri(n - 2, n - 3, -1));
    tripletList.push_back(Tri(n - 2, n - 2, 2));

    _A_basic.reserve(_maxdepth);
    SpM _A_temp(n - 1, n - 1);
    for (int i = 0; i < _maxdepth; i++)
    {
        _A_temp.resize(_width - 1, _width - 1);
        _A_temp.setFromTriplets(tripletList.begin(), tripletList.begin() + 3 * _width - 5);
        _A_basic.push_back(_A_temp);
        _width /= 2;
    };
    //
    // for (int i = 0; i < _maxdepth; i++)
    //     std::cout << _A_basic[i] << std::endl;
};

///尚未完工
void MultigridSolver::GenerateP(int n)
{
    int _maxdepth = _maxlevel + 1;
    int width = n;
    std::vector<Tri> tripletList1, tripletList2;
    tripletList1.reserve(3 * (n / 2 - 1));
    tripletList2.reserve(n / 2 - 1);
    for (int i = 0; i < n / 2 - 1; i++)
    {
        tripletList1.push_back(Tri(i, 2 * i, 1.0 / 4));
        tripletList1.push_back(Tri(i, 2 * i + 1, 2.0 / 4));
        tripletList1.push_back(Tri(i, 2 * i + 2, 1.0 / 4));
        tripletList2.push_back(Tri(i, 2 * i + 1, 1));
    };
    SpMbox _Restric_FW, _Intp_linear, _Restric_inj, _Intp_quad;
    _Restric_FW.reserve(_maxdepth);
    _Intp_linear.reserve(_maxdepth);
    _Restric_inj.reserve(_maxdepth);
    _Intp_quad.reserve(_maxdepth);

    /**
 * @brief 生成full weighting和injection.
 */
    SpM _SpM_temp(n / 2 - 1, n - 1);
    for (int i = 0; i < _maxdepth; i++)
    {
        _SpM_temp.resize(width / 2 - 1, width - 1);
        _SpM_temp.setFromTriplets(tripletList1.begin(), tripletList1.begin() + 3 * (width / 2 - 1));
        _Restric_FW.push_back(_SpM_temp);
        _SpM_temp.setZero();

        /**
         * @brief 对于vector容器类一种拙劣的初始化方式。
         */
        _Intp_linear.push_back(_SpM_temp);
        _Intp_quad.push_back(_SpM_temp);

        _SpM_temp.setFromTriplets(tripletList2.begin(), tripletList2.begin() + width / 2 - 1);
        _Restric_inj.push_back(_SpM_temp);
        width /= 2;
    };

    /**
 * @brief 通过full weighting生成intp linear。
 */
    for (int i = 0; i < _maxdepth; i++)
    {
        _Intp_linear[i] = _Restric_FW[i].transpose();
        _Intp_linear[i] = 2.0 * _Intp_linear[i];
    };

    switch (_Restric_method)
    {
    case full_weighting:
        _Restric.assign(_Restric_FW.begin(), _Restric_FW.end());
        break;

    case injection:
        _Restric.assign(_Restric_inj.begin(), _Restric_inj.end());
        break;
    default:

        break;
    };

    switch (_Intp_method)
    {
    case linear:
        _Intp.assign(_Intp_linear.begin(), _Intp_linear.end());
        break;
    default:
        break;
    };
};